from chiplet import QPU
from dag import Circuit_DAG
from expr_greedy import schedule_greedy
from expr_optimal import schedule_optimal
from expr_swap import schedule_swap

# circuit_types = ["bv"]
circuit_types = ["bv", "qaoa", "qft", "vqa"]
chip_types = ["square", "hexagon", "heavy_square", "heavy_hexagon"]
qubit_num = 100

for circuit_type in circuit_types:
    for chip_type in chip_types:
        qasm_file = f"circuits/{circuit_type}_{qubit_num}.qasm"
        circuit_dag = Circuit_DAG(qasm_file)
        # chip_type - Chip指向的类型，有 "square", "hexagon", "heavy_square", "heavy_hexagon" 四种
        is_square = circuit_type == "square"
        is_heavy_hexagon = circuit_type == "heavy_hexagon"
        qpu = QPU(chip_type=chip_type, inner_chip_size=(8 if is_square else (3 if is_heavy_hexagon else 5)), inter_chip_array=[
            3, 3], circuit_dag=circuit_dag, just_one_map=(circuit_type == "vqa" or circuit_type == "qft"))
        # max_depth, swap_num, ghz_length

        output_file_path = "output_chiplet_type_simple.txt"
        with open(output_file_path, 'a') as file:
            file.write(
                f"{qubit_num} {circuit_type} {chip_type}\n")

        print(f"{circuit_type}, {chip_type}")
        max_depth_swap = schedule_swap(qpu)
        print(f"max_depth_swap：{max_depth_swap}")
        with open(output_file_path, 'a') as file:
            file.write(
                f"{max_depth_swap}\n")

        max_depth_greedy = schedule_greedy(qpu)
        print(f"max_depth_greedy：{max_depth_greedy}")
        with open(output_file_path, 'a') as file:
            file.write(
                f"{max_depth_greedy}\n")

        # max_depth_optimal = schedule_optimal(qpu)
        # print(f"max_depth_optimal：{max_depth_optimal}")
        # with open(output_file_path, 'a') as file:
        #     file.write(
        #         f"{max_depth_optimal}\n")
